import os
from loguru import logger
from keras.preprocessing import image as image_utils
from keras.applications.imagenet_utils import decode_predictions
from keras.applications.imagenet_utils import preprocess_input
import numpy as np
from tensorflow.keras import applications
from tqdm import tqdm
import pickle as pk
import click
import test_model


@click.command()
@click.option("-m", "--model_type", default="Xception")
@click.option("-i", "--ori", default=0)
@click.option("-p", "--path")
def run(model_type, ori, path):

    IMAGE_SIZE = 224
    # path = "/data/gaohan/imagenet/labelled"
    filelist = os.listdir(path)
    logger.info("Image Count: {}".format(len(filelist)))

    model = None
    logger.info("Loadding Model ...")
    if ori == 0:
        # exec(open("test_model.py").read())
        model = test_model.model
    else:
        if model_type == "Xception":
            model = applications.Xception(weights='imagenet')
        elif model_type == "ResNet50":
            model = applications.ResNet50(weights='imagenet')
        elif model_type == "ResNet50V2":
            model = applications.ResNet50V2(weights="imagenet")
        elif model_type == "MobileNet":
            model = applications.MobileNet(weights='imagenet')
        elif model_type == "DenseNet121":
            model = applications.DenseNet121(weights="imagenet")
        else:
            logger.error("No Such Model!")
    if model_type == "Xception":
        IMAGE_SIZE = 299
    top5 = []
    gt = []
    for filename in tqdm(filelist):
        gt_id = filename.split(".")[-2]
        file_path = "{}/{}".format(path, filename)
        image = image_utils.load_img(file_path, target_size=(IMAGE_SIZE, IMAGE_SIZE))
        image = image_utils.img_to_array(image)
        image = np.expand_dims(image, axis=0)
        if model_type == "Xception":
            image = applications.xception.preprocess_input(image)
        elif model_type == "MobileNet":
            image = applications.mobilenet.preprocess_input(image)
        elif model_type == "ResNet50":
            image = applications.resnet.preprocess_input(image)
        elif model_type == "ResNet50V2":
            image = applications.resnet_v2.preprocess_input(image)
        elif model_type == "DenseNet121":
            image = applications.densenet.preprocess_input(image)
        else:
            image = preprocess_input(image)

        preds = model.predict(image)
        if len(preds.shape) == 3:
            preds = preds[0]
        P = decode_predictions(preds)
        tmp = []

        # logger.debug("Processing {}".format(filename))
        for (i, (imagenetID, label, prob)) in enumerate(P[0]):
            tmp.append(imagenetID)
            # logger.debug("{}. {}: {:.2f}%, ID[{}]".format(i + 1, label, prob * 100, imagenetID))
        top5.append(tmp)
        gt.append(gt_id)
        # break

    save_name = "{}-Result-{}.pkl".format(model_type, "Rev" if ori == 0 else "Ori")
    if os.path.exists(save_name):
        os.remove(save_name)

    with open(save_name, "wb") as f:
        pk.dump({"top5": top5, "gt": gt}, f)
    logger.info("Results are Saved!")


if __name__ == "__main__":
    run()
